#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2021/11/11 23:09
# @Author : 不想起名
# @Site :
# @File : OCR.py
# @Software: PyCharm
import os
from threading import Thread
import wx
from paddleocr import PaddleOCR
from wx import TextCtrl, adv

import screenShot

# PaddleOCR目前支持的多语言语种可以通过修改lang参数进行切换
# 例如`ch`, `en`, `fr`, `german`, `korean`, `japan`

ocr = PaddleOCR(det_model_dir='models/det/ch_PP-OCRv2_det_slim_quant_infer',
                rec_model_dir='models/rec/ch_PP-OCRv2_rec_slim_quant_infer',
                cls_model_dir='models/cls/ch_ppocr_mobile_v2.0_cls_slim_infer',
                use_angle_cls=True,
                lang="ch")  # need to run only once to download and load model into memory


class OCR_thread(Thread):
    """ 文本域类，子线程 """

    def __init__(self, parent, img_path):
        """
        :param parent:  主线程UI
        """
        super(OCR_thread, self).__init__()  # 继承
        self.parent = parent  # 主线程
        self.img_path = img_path  # 主线程传来的图片路径
        self.setDaemon(True)  # 设置为守护线程， 即子线程是守护进程，主线程结束子线程也随之结束。

    def run(self):
        result = ocr.ocr(self.img_path, cls=True)  # 在这另开线程执行ocr，不会导致主线程卡顿
        if result:
            for line in result:
                wx.CallAfter(self.parent.result_text.AppendText, line[1][0] + '\n')  # 主线程的文本域增加信息


def about_clicked(event):  # 这是MainWindow里面的关于
    description = """OCR是一款使用百度PaddlePaddle创建的跨平台的开源OCR文字识别软件"""
    developers = ['不想起名']

    info = wx.adv.AboutDialogInfo()
    info.SetName('OCR')
    info.SetIcon(wx.Icon('icons/OCR.png', wx.BITMAP_TYPE_PNG))
    info.SetDescription(description)
    info.SetCopyright('(C) 2021 - 2021 G-and-Z')
    info.SetWebSite('https://gitee.com/zhang-jia-liang12138/ocr', desc="OCR")
    info.SetVersion('1.0.0')
    info.SetDevelopers(developers)

    wx.adv.AboutBox(info)


class MainWindow(wx.Frame):

    def __init__(self, parent, title):
        super(MainWindow, self).__init__(parent, title=title)
        self.EnableMaximizeButton(False)  # 禁止最大化窗口
        self.SetMinSize((400, 520))
        self.SetMaxSize((400, 520))
        self.Centre()  # 将窗口显示在电脑系统中心
        self.SetIcon(wx.Icon("icons/OCR.png", wx.BITMAP_TYPE_PNG))  # Icon图标

        p = wx.Panel(self)

        # 垂直盒子
        vbox = wx.BoxSizer(wx.VERTICAL)
        # 水平盒子
        hbox = wx.BoxSizer(wx.HORIZONTAL)

        # 三个图标
        openfile = wx.Bitmap("icons/openfile.png", wx.BITMAP_TYPE_PNG)
        cut = wx.Bitmap("icons/cut.png", wx.BITMAP_TYPE_PNG)
        about = wx.Bitmap("icons/about.png", wx.BITMAP_TYPE_PNG)

        # 三个按钮
        but_cut = wx.BitmapButton(p, bitmap=cut)
        but_openfile = wx.BitmapButton(p, bitmap=openfile)
        but_about = wx.BitmapButton(p, bitmap=about)

        # 文本框
        font = wx.Font(12, wx.FONTFAMILY_DEFAULT, wx.FONTSTYLE_NORMAL, wx.FONTWEIGHT_NORMAL, False, 'Consolas')
        self.result_text = TextCtrl(p, style=wx.TE_MULTILINE, name="result_text")
        self.result_text.SetFont(font)

        # 三个按钮绑定的事件和函数
        but_cut.Bind(wx.EVT_BUTTON, lambda event, i=self.result_text: self.cut_clicked(event, i))
        but_openfile.Bind(wx.EVT_BUTTON, lambda event, i=self.result_text: self.openfile_clicked(event, i))
        but_about.Bind(wx.EVT_BUTTON, about_clicked)

        # Add第二个参数是在这个水平盒子里的比例，现在是1比1
        hbox.Add(but_cut, 1, wx.ALL | wx.EXPAND, 2)
        hbox.Add(but_openfile, 1, wx.ALL | wx.EXPAND, 2)
        hbox.Add(but_about, 1, wx.ALL | wx.EXPAND, 2)

        # Add第二个参数是在这个垂直盒子里的比例，现在是1比6
        vbox.Add(hbox, 1, wx.ALL | wx.EXPAND, 3)
        vbox.Add(self.result_text, 6, wx.ALL | wx.EXPAND, 3)

        p.SetSizer(vbox)
        self.Show()

    def cut_clicked(self, event, result_text):
        self.Iconize()  # 窗口最小化
        img_path = 'temp/tmp.png'
        if os.path.exists(img_path):
            os.unlink(img_path)  # 删除缓存图片
        screenShot.screenShot()  # 截图
        self.Iconize(False)  # 取消最小化
        if os.path.exists(img_path):  # 若截图得到新图片
            result_text.Clear()
            OCR_thread(self, img_path).start()  # 调用子线程

    def openfile_clicked(self, event, result_text):
        # 文件对话框
        with wx.FileDialog(self, "选择一张图片", wildcard="BMP、PNG、JPG、JPEG files ("
                                                    "*.bmp;*.png;*.jpg;*.jpeg)|*.bmp;*.png;*.jpg;*.jpeg",
                           style=wx.FD_OPEN | wx.FD_FILE_MUST_EXIST) as fileDialog:
            if fileDialog.ShowModal() == wx.ID_CANCEL:
                return  # 取消选择
            # 选取之后执行以下代码
            pathname = fileDialog.GetPath()
            result_text.Clear()
            OCR_thread(self, pathname).start()


app = wx.App()
MainWindow(None, title='OCR')
app.MainLoop()
